package com.mkx.orderedlist.avl;

import com.mkx.orderedlist.OrderedList;

/**
 * 基于AVL树实现的有序表
 * @param <K> 键类型
 * @param <V> 值类型
 */
public class AVLTreeOrderedList<K extends Comparable<K>, V> implements OrderedList<K, V> {
    private AVLTreeNode<K, V> root;
    private int size;
    
    /**
     * 构造一个空的AVL树有序表
     */
    public AVLTreeOrderedList() {
        root = null;
        // 明确初始化size为0
        this.size = 0;
    }
    
    /**
     * 获取节点的高度
     * @param node 节点
     * @return 高度值
     */
    private int height(AVLTreeNode<K, V> node) {
        if (node == null) {
            return 0;
        }
        return node.height;
    }
    
    /**
     * 计算平衡因子
     * @param node 节点
     * @return 平衡因子（左子树高度减去右子树高度）
     */
    private int balanceFactor(AVLTreeNode<K, V> node) {
        if (node == null) {
            return 0;
        }
        return height(node.left) - height(node.right);
    }
    
    /**
     * 更新节点的高度
     * @param node 节点
     */
    private void updateHeight(AVLTreeNode<K, V> node) {
        if (node != null) {
            node.height = Math.max(height(node.left), height(node.right)) + 1;
        }
    }
    
    /**
     * 右旋转操作
     * @param y 旋转中心
     * @return 新的根节点
     */
    private AVLTreeNode<K, V> rightRotate(AVLTreeNode<K, V> y) {
        AVLTreeNode<K, V> x = y.left;
        AVLTreeNode<K, V> T3 = x.right;
        
        // 执行旋转
        x.right = y;
        y.left = T3;
        
        // 更新高度
        updateHeight(y);
        updateHeight(x);
        
        return x;
    }
    
    /**
     * 左旋转操作
     * @param x 旋转中心
     * @return 新的根节点
     */
    private AVLTreeNode<K, V> leftRotate(AVLTreeNode<K, V> x) {
        AVLTreeNode<K, V> y = x.right;
        AVLTreeNode<K, V> T2 = y.left;
        
        // 执行旋转
        y.left = x;
        x.right = T2;
        
        // 更新高度
        updateHeight(x);
        updateHeight(y);
        
        return y;
    }
    
    /**
     * 平衡节点
     * @param node 需要平衡的节点
     * @return 平衡后的节点
     */
    private AVLTreeNode<K, V> balance(AVLTreeNode<K, V> node) {
        if (node == null) {
            return null;
        }
        
        // 更新当前节点的高度
        updateHeight(node);
        
        // 获取平衡因子
        int balance = balanceFactor(node);
        
        // 左侧不平衡情况
        if (balance > 1) {
            // 左-左情况
            if (balanceFactor(node.left) >= 0) {
                return rightRotate(node);
            }
            // 左-右情况
            node.left = leftRotate(node.left);
            return rightRotate(node);
        }
        
        // 右侧不平衡情况
        if (balance < -1) {
            // 右-右情况
            if (balanceFactor(node.right) <= 0) {
                return leftRotate(node);
            }
            // 右-左情况
            node.right = rightRotate(node.right);
            return leftRotate(node);
        }
        
        return node; // 已经平衡
    }
    
    /**
     * 插入节点的递归辅助方法
     * @param node 当前节点
     * @param key 键
     * @param value 值
     * @return 插入后的新根节点
     */
    @Override
    public V put(K key, V value) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        
        // 保存旧值
        V oldValue = get(key);
        
        // 执行插入操作
        root = putRecursive(root, key, value);
        
        // 如果键不存在，增加size
        if (oldValue == null) {
            size++;
        }
        
        return oldValue;
    }
    
    private AVLTreeNode<K, V> putRecursive(AVLTreeNode<K, V> node, K key, V value) {
        // 执行标准的BST插入
        if (node == null) {
            return new AVLTreeNode<>(key, value);
        }
        
        int compareResult = key.compareTo(node.key);
        
        if (compareResult < 0) {
            node.left = putRecursive(node.left, key, value);
        } else if (compareResult > 0) {
            node.right = putRecursive(node.right, key, value);
        } else {
            // 键已存在，更新值
            node.value = value;
            return node; // 值更新不需要旋转
        }
        
        // 平衡树
        return balance(node);
    }
    
    /**
     * 查找节点的递归辅助方法
     * @param node 当前节点
     * @param key 键
     * @return 对应的值
     */
    private V getRecursive(AVLTreeNode<K, V> node, K key) {
        if (node == null) {
            return null;
        }
        
        int compareResult = key.compareTo(node.key);
        
        if (compareResult < 0) {
            return getRecursive(node.left, key);
        } else if (compareResult > 0) {
            return getRecursive(node.right, key);
        } else {
            return node.value;
        }
    }
    
    @Override
    public V get(K key) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        return getRecursive(root, key);
    }
    
    /**
     * 查找最小键的节点
     * @param node 当前节点
     * @return 最小键的节点
     */
    private AVLTreeNode<K, V> findMin(AVLTreeNode<K, V> node) {
        AVLTreeNode<K, V> current = node;
        while (current.left != null) {
            current = current.left;
        }
        return current;
    }
    
    /**
     * 删除节点的递归辅助方法
     * @param node 当前节点
     * @param key 键
     * @return 删除后的新根节点
     */
    private AVLTreeNode<K, V> removeRecursive(AVLTreeNode<K, V> node, K key) {
        if (node == null) {
            return null;
        }
        
        int compareResult = key.compareTo(node.key);
        
        if (compareResult < 0) {
            node.left = removeRecursive(node.left, key);
        } else if (compareResult > 0) {
            node.right = removeRecursive(node.right, key);
        } else {
            // 找到要删除的节点
            size--;
            
            // 情况1: 叶节点
            // 情况2: 只有一个子节点
            if (node.left == null) {
                return node.right;
            } else if (node.right == null) {
                return node.left;
            }
            
            // 情况3: 有两个子节点
            // 找到右子树中的最小节点（后继节点）
            AVLTreeNode<K, V> successor = findMin(node.right);
            
            // 替换当前节点的值
            node.key = successor.key;
            node.value = successor.value;
            
            // 删除后继节点
            node.right = removeRecursive(node.right, successor.key);
        }
        
        // 平衡树
        return balance(node);
    }
    
    @Override
    public V remove(K key) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        
        // 保存旧值
        V oldValue = get(key);
        
        if (oldValue != null) {
            // 执行删除操作
            root = removeRecursive(root, key);
            
            // 确保size减少
            size = Math.max(0, size - 1);
        }
        
        return oldValue;
    }
    
    @Override
    public boolean containsKey(K key) {
        return get(key) != null;
    }
    
    @Override
    public int size() {
        // 使用递归方法计算实际大小，确保准确性
        return calculateSize(root);
    }
    
    /**
     * 递归计算树的大小
     * @param node 当前节点
     * @return 以当前节点为根的子树大小
     */
    private int calculateSize(AVLTreeNode<K, V> node) {
        if (node == null) {
            return 0;
        }
        return 1 + calculateSize(node.left) + calculateSize(node.right);
    }
    
    @Override
    public boolean isEmpty() {
        // 根据实际计算的大小判断是否为空
        return size() == 0;
    }
    
    @Override
    public void clear() {
        root = null;
        size = 0;
    }
    
    @Override
    public K firstKey() {
        if (root == null) {
            return null;
        }
        AVLTreeNode<K, V> current = root;
        while (current.left != null) {
            current = current.left;
        }
        return current.key;
    }
    
    @Override
    public K lastKey() {
        if (root == null) {
            return null;
        }
        AVLTreeNode<K, V> current = root;
        while (current.right != null) {
            current = current.right;
        }
        return current.key;
    }
}